import chromadb
from langchain_chroma import Chroma as LangchainChroma
import sys
import os
user_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'customize')
# 将 user 目录添加到 sys.path 中
sys.path.append(user_dir)
db_dir = os.path.join(os.path.dirname(__file__), 'chroma_db')
from customize.embedding import get_embedding

class ChromaManager:

    def __init__(self, collection_name="knowledge_base", persist_dir=db_dir):
        self.client = chromadb.PersistentClient(path=persist_dir)
        self.embedding_func = get_embedding(
            model_path="F:/models/BAAI/bge-large-zh-v1.5",
            device='cuda',
            normalize_embeddings=True
        )
        self.langchain_chroma = LangchainChroma(
            client=self.client,
            collection_name=collection_name,
            embedding_function=self.embedding_func
        )

    def add_document(self, doc_id: str, content: str, metadata: dict=None):
        """添加/更新文档"""
        self.langchain_chroma.add_texts(
            texts=[content],
            ids=[doc_id],
            metadatas=[metadata] if metadata else None
        )

    def delete_document(self, doc_id: str):
        """删除文档"""
        self.langchain_chroma.delete(ids=[doc_id])

    def search(self, query: str, n_results: int=5):
        """搜索文档"""
        return self.langchain_chroma.similarity_search_with_score(query, k=n_results)

    def get_all_documents(self):
        """获取全部文档（简单实现）"""
        return self.langchain_chroma.get()
